#我现在有水位尺的图片和一些水的图片，需要将水的图片和水位尺的图片混合；并导出xml标注文件
import argparse
from pathlib import Path
import cv2
import numpy as np
import random
from PIL import Image
import xml.etree.ElementTree as ET

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument('--background-dir', default=r"D:\data\研究生\课题资料\水位尺\red_line", type=str)
    parser.add_argument('--fore-dir', default=r"D:\data\研究生\课题资料\水位尺\water", type=str)
    parser.add_argument('--output-img-dir', default=r"D:\data\研究生\课题资料\水位尺\images", type=str)
    parser.add_argument('--output-xml-dir', default=r"D:\data\研究生\课题资料\水位尺\annotations", type=str)
    return parser.parse_args()

def read_image_from_chinese_path(chinese_path):
    # 使用numpy从中文路径读取数据
    file_stream = open(chinese_path, 'rb')
    file_bytes = np.asarray(bytearray(file_stream.read()), dtype=np.uint8)
    file_stream.close()

    # 使用cv2.imdecode读取图片数据
    img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)

    return img
#获取图片中的区域
def getRandomRegion(img, w, h):
    try:
        img_height, img_width = img.shape[:2]
        if img_height < h or img_width < w:
            img = cv2.resize(img, (w,h))
            img_height, img_width = img.shape[:2]
        x_start = np.random.randint(0, img_width - w + 1)
        y_start = np.random.randint(0, img_height - h + 1)
        cropped_img = img[y_start:y_start+h, x_start:x_start+w]
        return cropped_img
    except Exception as e:
        print("裁剪区域失败", e)
        return None
#美化xml
def prettyXml(element, indent, newline, level = 0): # elemnt为传进来的Elment类，参数indent用于缩进，newline用于换行
    if element:  # 判断element是否有子元素
        if element.text == None or element.text.isspace(): # 如果element的text没有内容
            element.text = newline + indent * (level + 1)
        else:
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)
    #else:  # 此处两行如果把注释去掉，Element的text也会另起一行
        #element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level
    temp = list(element) # 将elemnt转成list
    for subelement in temp:
        if temp.index(subelement) < (len(temp) - 1): # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致
            subelement.tail = newline + indent * (level + 1)
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个
            subelement.tail = newline + indent * level
        prettyXml(subelement, indent, newline, level = level + 1) # 对子元素进行递归操作
#生成voc_xml文件
def create_voc_xml(filename, width, height, depth, object_name, bbox):
    root = ET.Element("annotation")

    folder = ET.SubElement(root, "folder")
    folder.text = "images"

    filename_xml = ET.SubElement(root, "filename")
    filename_xml.text = filename

    path_xml = ET.SubElement(root, "path")
    path_xml.text = filename

    source = ET.SubElement(root, "source")
    database = ET.SubElement(source, "database")
    database.text = "Unknown"

    size = ET.SubElement(root, "size")
    width_xml = ET.SubElement(size, "width")
    width_xml.text = str(width)
    height_xml = ET.SubElement(size, "height")
    height_xml.text = str(height)
    depth_xml = ET.SubElement(size, "depth")
    depth_xml.text = str(depth)

    segmented = ET.SubElement(root, "segmented")
    segmented.text = "0"

    object_xml = ET.SubElement(root, "object")
    name = ET.SubElement(object_xml, "name")
    name.text = object_name
    pose = ET.SubElement(object_xml, "pose")
    pose.text = "Unspecified"
    truncated = ET.SubElement(object_xml, "truncated")
    truncated.text = "0"
    difficult = ET.SubElement(object_xml, "difficult")
    difficult.text = "0"
    bndbox = ET.SubElement(object_xml, "bndbox")
    xmin = ET.SubElement(bndbox, "xmin")
    xmin.text = str(bbox[0])
    ymin = ET.SubElement(bndbox, "ymin")
    ymin.text = str(bbox[1])
    xmax = ET.SubElement(bndbox, "xmax")
    xmax.text = str(bbox[2])
    ymax = ET.SubElement(bndbox, "ymax")
    ymax.text = str(bbox[3])

    # 创建ElementTree对象并保存为XML文件
    prettyXml(root, '    ', '\n')  #执行美化方法
    tree = ET.ElementTree(root)
    return tree

def main():
    args = Parse_Arguments()
    background_img_dir = list(Path(args.background_dir).glob("*.jpg"))
    fore_dir = list(Path(args.fore_dir).glob("*.jpg"))
    output_img_dir = args.output_img_dir
    output_xml_dir = args.output_xml_dir
    output_img_cnt = 300


    for _ in range(500):
        for background_path in background_img_dir:
            for fore_path in fore_dir:
                mix_ratios = 0 #记录水占标靶的像素值
                alpha = 0.2
                background_img = read_image_from_chinese_path(str(background_path))  #读取中文路径下的图片
                fore_img = read_image_from_chinese_path(str(fore_path))
                background_img_height,background_img_width= background_img.shape[:2]
                region_height = int(random.uniform(5, 10)) #随机裁剪水图片的高度占原图高度的比例
                while mix_ratios < 80:
                    fore_region = getRandomRegion(fore_img, background_img_width, region_height) #随机获取水中的区域
                    mix_ratios += region_height
                    mix_leftTop_position_y = background_img_height - mix_ratios #叠在原图的左上角y坐标
                    background_region = background_img[mix_leftTop_position_y:mix_leftTop_position_y + region_height, 0:background_img_width] #取出背景图片一部分
                    mix_region = cv2.addWeighted(background_region, alpha, fore_region, 1-alpha, 0)  #混合图片
                    background_img[mix_leftTop_position_y:mix_leftTop_position_y + region_height, 0:background_img_width] = mix_region #将混合图片放回原图
                    save_img = cv2.cvtColor(background_img, cv2.COLOR_BGR2RGB)  # PIL需要RGB格式
                    save_img = Image.fromarray(save_img)
                    output_filename_img = f"{output_img_dir}\\{output_img_cnt:04d}.jpg"
                    output_filename_xml = f"{output_xml_dir}\\{output_img_cnt:04d}.xml"
                    filename_img = f"{output_img_cnt:04d}.jpg"
                    xml_tree = create_voc_xml(filename_img,width=72,height=mix_leftTop_position_y - 36,depth=3,object_name="red_line",bbox=[35,36,107,mix_leftTop_position_y])
                    xml_tree.write(output_filename_xml)
                    save_img.save(output_filename_img)
                    output_img_cnt +=1
                    if output_img_cnt > 400:
                        exit()


if __name__ == "__main__":
    main()
